/*
 * bcd_sequential.cc
 *
 * Copyright 2013 Long Tran (ltran3@gatech.edu)
 *
 * BlockCoodinateDescent, solving for local minimum of
 *   L = sum |A - W H^T|^2
 *   subject to W >= 0, H >= 0
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

#include <armadillo>
#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/mpi.hpp>

#include "helper/helper.h"
#include "algorithm/StopWatch.h"
#include "data/FileAccess.h"
extern boost::mpi::communicator* p_world;
extern po::variables_map vm;
#include "mpi_helper.h"
#include "arma_helper.h"

class BlockCoodinateDescent
{
public:
  // input/output variables
  arma::mat A, W, H;
  // iteration information
  std::vector<Iteration> iterInfo;
protected:
  // algorithm variables
  int m, n, k;
  double l1, l2;
  int maxiter, interval;
  bool normalizing;
public:
  BlockCoodinateDescent() {
    PO_REQUIRE(vm, "iter");
    PO_REQUIRE(vm, "k");

    PO_VARIABLE_VALUE(k, vm, "k", int);
    l1 = (vm.count("l1")) ? vm["l1"].as<double>() : 0; // l1 regularization
    l2 = (vm.count("l2")) ? vm["l2"].as<double>() : 0; // l2 regularization
    PO_VARIABLE_VALUE(maxiter, vm, "iter", int);
    PO_VARIABLE_VALUE(interval, vm, "interval", int);
    normalizing = vm.count("normalize");
  }

  void LoadData() 
  {
    PO_REQUIRE(vm, "data");
    PO_VARIABLE(data, vm, "data", std::string);
    
    using namespace arma;
    std::stringstream ss;
    ss << data << "_" << MPI_RANK;
    A.load(ss.str(), arma_binary);
    m = A.n_rows;
    n = A.n_cols;
  }

  static void ThresholdVector(arma::vec& v)
  {
    for (unsigned int i = 0; i < v.n_elem; i++)
      if (v(i) < 1e-16) v(i) = 0;
  }

  void SolveBCD()
  {
    StopWatch watch, clock;
    TimeDuration runtime;

    using namespace arma;
    srand(100);
    W = randu<mat>(m, k); 
    H = randu<mat>(n, k); 

    for (int iter = 0; iter < maxiter; iter++) {
      clock.GetDuration();
      // update W
      mat HH = trans(H)*H;
      mat AH = A*H;
      for (int i = 0; i < k; i++) {
	double hii = HH(i,i);
	arma::vec w_i = (AH.col(i) - W*HH.col(i) + hii*W.col(i)) / hii;
	ThresholdVector(w_i);
	W.col(i) = w_i;
      }

      // update H
      mat WW = trans(W)*W;
      mat AW = trans(A)*W;
      for (int i = 0; i < k; i++) {
	double wii = WW(i,i);
	arma::vec h_i = ( AW.col(i) - H*WW.col(i) + wii*H.col(i) ) / wii;
	ThresholdVector(h_i);
	H.col(i) = h_i;
      }
      runtime += clock.GetDuration();

      if (iter % interval == 0 || iter == maxiter-1) {
        double res = norm(A-W*trans(H), "fro");
        if (MPI_RANK == 0) {
          MSG("it=" << iter << " |res|=" << res  
	      << " |W|=" << norm(W,"fro") 
	      << " |H|=" << norm(H,"fro")
              << " d=" << watch.GetDuration());
	  Iteration info(iter, res);
	  info.duration = watch.GetTotalDuration();
	  info.compute = runtime;
	  iterInfo.push_back(info);
	}
      }
    }
    MSG(watch.GetTotalDuration());
      MSG("compute clock " << runtime);
    
  }
};

void run_bcd()
{
  if (MPI_RANK == 0) MSG("Block Coordinate Descent NMF");
  BlockCoodinateDescent a;
  a.LoadData();
  a.SolveBCD();
  
  if (MPI_RANK == 0 && vm.count("output")) {
    PO_VARIABLE(output, vm, "output", std::string);
    a.H.save(output, arma::raw_ascii);
  }

  if (MPI_RANK == 0 && vm.count("iterinfo")) {
    PO_VARIABLE(infoFile, vm, "iterinfo", std::string);
    FileAccessWrite(out, infoFile, NULL);
    const std::vector<Iteration>& info =  a.iterInfo;
    for (unsigned int i = 0; i < info.size(); i++) {
      out << info[i].iter << "\t" << (double)info[i].duration.total_milliseconds()/1000 << "\t" 
	  << (double)info[i].compute.total_milliseconds()/1000 << "\t" 
	  << (double)info[i].communication.total_milliseconds()/1000 << "\t"
	  << info[i].residual << std::endl;
    }
  }
}
